import * as tf from "@tensorflow/tfjs";
import * as tfvis from "@tensorflow/tfjs-vis";
import { getIrisData, IRIS_CLASSES } from "./data";

(async () => {
    /**
     * [10,2,0,10]
     * [[1,0,0]]
     * ["xx花","yy花","zz花"]
     * 
     */
  let [xTrain, yTrain, xTest, yTest] = getIrisData(0.01);

  let model = tf.sequential();

  model.add(
    tf.layers.dense({
      units: 10,
      inputShape: [4],
      activation: "sigmoid",
    })
  );

  model.add(
    tf.layers.dense({
      units: 3, // 输出的神经元
      activation: "softmax", // 输出三个概率
    })
  );

  model.compile({
    loss: "categoricalCrossentropy",
    optimizer: tf.train.adam(0.01),
    metrics: ["accuracy"], // 度量单位 accuracy准确度数
  });
  // 获取表单的提交
  window["predict"] = ()=>{
    let formData = document.form;
    let input = tf.tensor([[
        formData.a.value * 1,
        formData.b.value * 1,
        formData.c.value * 1,
        formData.d.value * 1,
    ]])
    let pred =  model.predict(input); // [ [] ]
    console.log(pred.print());
    console.log(IRIS_CLASSES[pred.argMax(1).dataSync(0)]);
    return false;
  }

  await model.fit(xTrain, yTrain, {
    epochs: 100,
    validationData: [xTest, yTest],
    callbacks: tfvis.show.fitCallbacks(
      {
        name: "训练过程",
      },
      ["loss", "val_loss", "acc", "val_acc"],
      {
        callbacks: ["onEpochEnd"],
      }
    ),
  });
  console.log("训练完成");
})();
